# RUN: tf-mlir-translate -graphdef-to-mlir -tf-upgrade-legacy %s -tf-output-arrays=hash_table_node,variable_node,variable_v2_node -o - | FileCheck %s

node: {
  name: "hash_table_node"
  op: "HashTableV2"
  attr: {
    key: "key_dtype"
    value: {
      type: DT_INT32
    }
  }
  attr: {
    key: "shared_name"
    value: {
      s: ""
    }
  }
  attr: {
    key: "value_dtype"
    value: {
      type: DT_FLOAT
    }
  }
}
node {
  name: "variable_node"
  op: "Variable"
  attr {
    key: "dtype"
    value {
      type: DT_INT64
    }
  }
  attr {
    key: "shape"
    value {
      shape {
      }
    }
  }
  attr {
    key: "shared_name"
    value {
      s: ""
    }
  }
}
node {
  name: "variable_v2_node"
  op: "VariableV2"
  attr {
    key: "dtype"
    value {
      type: DT_INT64
    }
  }
  attr {
    key: "shape"
    value {
      shape {
      }
    }
  }
  attr {
    key: "shared_name"
    value {
      s: ""
    }
  }
}
node {
  name: "Call"
  op: "PartitionedCall"
  attr {
    key: "Tin"
    value {
      list {
      }
    }
  }
  attr {
    key: "Tout"
    value {
      list {
        type: DT_RESOURCE
      }
    }
  }
  attr {
    key: "f"
    value {
      func {
        name: "create_resource"
      }
    }
  }
}
node: {
  name: "batch_node"
  op: "BatchFunction"
  input: ["variable_node", "variable_v2_node"]
  attr: {
    key: "Tcaptured"
    value: {
      list: {
        type: [DT_INT64]
      }
    }
  }
  attr: {
    key: "Tin"
    value: {
      list: {
        type: [DT_INT64]
      }
    }
  }
  attr: {
    key: "Tout"
    value: {
      list: {
        type: [DT_INT64]
      }
    }
  }
  attr: {
    key: "batch_timeout_micros"
    value: {
      i: 100000
    }
  }
  attr: {
    key: "f"
    value: {
      func: {
        name: "inference_function"
      }
    }
  }
  attr: {
    key: "max_batch_size"
    value: {
      i: 32
    }
  }
  attr: {
    key: "num_batch_threads"
    value: {
      i: 16
    }
  }
  attr: {
    key: "shared_name"
    value: {
      s: ""
    }
  }
}
library {
  function {
    signature {
      name: "create_resource"
      output_arg {
        name: "handle"
        type: DT_RESOURCE
      }
    }
    node_def: {
      name: "hash_table_node"
      op: "HashTableV2"
      attr: {
        key: "key_dtype"
        value: {
          type: DT_INT32
        }
      }
      attr: {
        key: "shared_name"
        value: {
          s: ""
        }
      }
      attr: {
        key: "value_dtype"
        value: {
          type: DT_FLOAT
        }
      }
    }
    node_def {
      name: "variable_node"
      op: "Variable"
      attr {
        key: "dtype"
        value {
          type: DT_INT64
        }
      }
      attr {
        key: "shape"
        value {
          shape {
          }
        }
      }
      attr {
        key: "shared_name"
        value {
          s: ""
        }
      }
    }
    node_def {
      name: "variable_v2_node"
      op: "VariableV2"
      attr {
        key: "dtype"
        value {
          type: DT_INT64
        }
      }
      attr {
        key: "shape"
        value {
          shape {
          }
        }
      }
      attr {
        key: "shared_name"
        value {
          s: ""
        }
      }
    }
    node_def {
      name: "batch_node"
      op: "BatchFunction"
      input: ["variable_node:ref", "variable_v2_node:ref"]
      attr: {
        key: "Tcaptured"
        value: {
          list: {
            type: [DT_INT64]
          }
        }
      }
      attr: {
        key: "Tin"
        value: {
          list: {
            type: [DT_INT64]
          }
        }
      }
      attr: {
        key: "Tout"
        value: {
          list: {
            type: [DT_INT64]
          }
        }
      }
      attr: {
        key: "batch_timeout_micros"
        value: {
          i: 100000
        }
      }
      attr: {
        key: "f"
        value: {
          func: {
            name: "inference_function"
          }
        }
      }
      attr: {
        key: "max_batch_size"
        value: {
          i: 32
        }
      }
      attr: {
        key: "num_batch_threads"
        value: {
          i: 16
        }
      }
      attr: {
        key: "shared_name"
        value: {
          s: ""
        }
      }
    }
    ret {
      key: "handle"
      value: "hash_table_node:table_handle:0"
    }
  }
  function {
    signature {
      name: "inference_function"
      output_arg {
        name: "handle"
        type: DT_RESOURCE
      }
    }
    node_def: {
      name: "hash_table_node"
      op: "HashTableV2"
      attr: {
        key: "key_dtype"
        value: {
          type: DT_INT32
        }
      }
      attr: {
        key: "shared_name"
        value: {
          s: ""
        }
      }
      attr: {
        key: "value_dtype"
        value: {
          type: DT_FLOAT
        }
      }
    }
    ret {
      key: "handle"
      value: "hash_table_node:table_handle:0"
    }
  }
}

# CHECK: tf.HashTableV2
# CHECK-SAME: shared_name = "hash_table_node"

# CHECK: "tf.Variable"
# CHECK-SAME: shared_name = "variable_node"
# CHECK: "tf.VariableV2"
# CHECK-SAME: shared_name = "variable_v2_node"

# CHECK: "tf.BatchFunction"
# CHECK-SAME: shared_name = "batch_node"

# CHECK: func private @create_resource
# CHECK: tf.HashTableV2
# CHECK-SAME: shared_name = "hash_table_node@create_resource"
# CHECK: "tf.Variable"
# CHECK-SAME: shared_name = "variable_node"
# CHECK: "tf.VariableV2"
# CHECK-SAME: shared_name = "variable_v2_node"
# CHECK: "tf.BatchFunction"
# CHECK-SAME: shared_name = "batch_node"
